import argparse
import numpy as np
from models.subchannel import SubChannel
from solvers.mass_momentum_energy import TwoFluidSolver
from solvers.convergence import ConvergenceChecker
from utils.plotting import Plotter
from utils.data_processing import DataProcessor
from utils.data_logger import DataLogger
import config

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='Two-fluid subchannel analysis program')
    parser.add_argument('--n_timesteps', type=int, default=None,
                       help='Number of time steps')
    parser.add_argument('--save_results', action='store_true',
                       help='Save calculation results')
    parser.add_argument('--plot_results', action='store_true',
                       help='Plot calculation results')
    parser.add_argument('--output_prefix', type=str, default='results',
                       help='Prefix for output files')
    return parser.parse_args()

def main():
    # 解析命令行参数
    args = parse_arguments()
    
    try:
        # 初始化数据记录器
        logger = DataLogger()
        
        # 记录计算参数
        params = {
            'n_timesteps': args.n_timesteps if args.n_timesteps is not None else int(config.NUMERICAL_PARAMS['total_time'] / config.NUMERICAL_PARAMS['time_step']),
            'n_axial': config.GRID_PARAMS['n_axial'],
            'n_channels': config.GRID_PARAMS['n_channels'],
            'time_step': config.NUMERICAL_PARAMS['time_step'],
            'convergence_tolerance': config.NUMERICAL_PARAMS['convergence_tolerance']
        }
        logger.log_calculation_parameters(params)
        
        # 初始化子通道
        subchannel = SubChannel()
        
        # 初始化求解器
        solver = TwoFluidSolver(subchannel)
        convergence_checker = ConvergenceChecker(config.NUMERICAL_PARAMS['convergence_tolerance'])
        
        # 开始瞬态计算
        solver.solve_transient()
        
        # 如果需要保存结果
        if args.save_results:
            processor = DataProcessor(subchannel)
            processor.export_results_to_excel(f'{args.output_prefix}_results.xlsx')
            stats = processor.generate_summary_statistics()
            logger.save_results(stats)
        
        # 如果需要绘图
        if args.plot_results:
            plotter = Plotter(subchannel)
            # 绘制空泡率轴向分布
            plotter.plot_axial_distribution(
                subchannel.void_fraction,
                'Axial Void Fraction Distribution',
                'Void Fraction'
            )
            # 绘制压力轴向分布
            plotter.plot_axial_distribution(
                subchannel.pressure,
                'Axial Pressure Distribution',
                'Pressure (Pa)'
            )
            # 绘制空泡率3D分布
            plotter.plot_void_fraction_3d()
            # 绘制收敛历史
            plotter.plot_convergence_history(
                convergence_checker.mass_balance_history,
                convergence_checker.energy_balance_history
            )
            
    except Exception as e:
        print(f"Error during calculation: {e}")
        import traceback
        traceback.print_exc()
        return False
        
    return True

if __name__ == '__main__':
    main() 